Status

• 🚧 NEED pytorch/pytorch changes 🚧
  • [c10d] support dynamic shapes for all_to_all_single_autograd pytorch#157521
  • Fix grouped MM output strides when compiled but not max-autotuned pytorch#158143
• tlparse: https://fburl.com/6g9djppf
  • staticmethod on user-defined classes can not be generically supported, I moved those out.
  • Remaining graph breaks are due to:
    • MoE experts being wrapped by FSDP.
      • FSDP hooks are explicitly disabled, we are graph breaking when the nn.Module.call invokes the forward pre and post hooks.
      • Since we will not be tracing FSDP wrapper, these graph breaks are considered fundamental until we migrate to SimpleFSDP.
    • Not technically a graph break, but the EP DTensor hooks are being traced in their own graphs. Still needs investigation.

We don't have a good way in compile to specify fullgraph=True except for FSDP hooks at the moment. We can either leave it fullgraph=False or just wrap the experts model code in set_fullgraph(False)/set_fullgraph(True).

Repro

tested on debug model NGPU=2 CONFIG_FILE="./torchtitan/experiments/llama4/train_configs/debug_model.toml" ./run_train.sh --parallelism.data_parallel_shard_degree=2 --parallelism.expert_parallel_degree=2 --training.compile
logs: https://gist.github.com/xmfan/41b822d9f09eb07fee62d684a061cec1

memory: 2.20GiB -> 1.42GiB
speedup: no big change, need to check with actual model